import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from matplotlib.font_manager import FontProperties
import pywt
from scipy.fft import fft, fftfreq

try:
    font = FontProperties(fname=r"C:\Windows\Fonts\msyh.ttc")
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei']
    plt.rcParams['axes.unicode_minus'] = False
except:
    plt.rcParams['font.sans-serif'] = ['SimHei', 'KaiTi', 'Arial Unicode MS']
    plt.rcParams['axes.unicode_minus'] = False

matplotlib.use('TkAgg', force=True)

sample_signal = np.load(r"D:\A_90762\atl\Voltage\train\class1\D52531318073.npy") # 取一个信号

sample_signal = sample_signal[::1]

fs = 1e6

fig, axes = plt.subplots(3, 1, figsize=(12, 12))

time = np.arange(len(sample_signal)) / fs
total_time = len(sample_signal) / fs

# 1. 绘制时域图
axes[0].plot(time, sample_signal)
axes[0].set_title('时域图')
axes[0].set_xlabel('时间 (秒)')
axes[0].set_ylabel('幅度')
axes[0].grid(True)
axes[0].set_xlim(0, total_time)  # 设置x轴范围与总时间一致

# 2. 绘制频域图
fft_result = fft(sample_signal)
fft_magnitude = np.abs(fft_result)
frequencies = fftfreq(len(sample_signal), 1/fs)

# 只取正频率部分
positive_freq_mask = frequencies >= 0
frequencies_positive = frequencies[positive_freq_mask]
fft_magnitude_positive = fft_magnitude[positive_freq_mask]

axes[1].plot(frequencies_positive, fft_magnitude_positive)
axes[1].set_title('频域图')
axes[1].set_xlabel('频率 (Hz)')
axes[1].set_ylabel('幅度')
axes[1].grid(True)
axes[1].set_xlim(0, fs/2)  # 只显示奈奎斯特频率以下的部分

# 3. 绘制时频图 (CWT)
scales = np.arange(1, 128) #第一个数决定最高频率，第二个数决定最低频率。例如：（1,32）的图几乎只有高频；(64,128)的图几乎只有低频
coef, freqs = pywt.cwt(sample_signal, scales, 'morl')

# 计算频率轴（将尺度转换为近似频率）
# 对于Morlet小波，频率 ≈ 中心频率 / (尺度 * 采样间隔)
center_frequency = pywt.central_frequency('morl')  # 获取Morlet小波的中心频率
frequencies_cwt = center_frequency * fs / scales  # 将尺度转换为频率

# 绘制时频图
im = axes[2].imshow(np.abs(coef), extent=[0, total_time, frequencies_cwt[-1], frequencies_cwt[0]],
           aspect='auto', cmap='jet',interpolation='nearest')
axes[2].set_title('CWT时频图')
axes[2].set_ylabel('频率 (Hz)')
axes[2].set_xlabel('时间 (秒)')
axes[2].set_xlim(0, total_time)
# plt.colorbar(im, ax=axes[2], label='系数幅度')

plt.tight_layout()
plt.subplots_adjust(hspace=0.4)
plt.show()

